import trimesh
import numpy as np
import torch
import pytorch3d
import pytorch3d.io as IO
from pytorch3d.transforms import quaternion_apply, matrix_to_quaternion

def get_bounding_box_coords(mesh_verts, bounding_box):
    '''
    :param mesh: trimesh mesh
    :return: 8 x 3
    '''
    center = torch.mean(mesh_verts, dim=0).numpy()
    print(center)
    bounding_box = bounding_box.numpy()
    coords = []
    coords.append([center[0] + bounding_box[0][1], center[1] + bounding_box[1][1], center[2] + bounding_box[2][1]])
    coords.append([center[0] + bounding_box[0][1], center[1] + bounding_box[1][1], center[2] + bounding_box[2][0]])
    coords.append([center[0] + bounding_box[0][1], center[1] + bounding_box[1][0], center[2] + bounding_box[2][1]])
    coords.append([center[0] + bounding_box[0][1], center[1] + bounding_box[1][0], center[2] + bounding_box[2][0]])
    coords.append([center[0] + bounding_box[0][0], center[1] + bounding_box[1][1], center[2] + bounding_box[2][1]])
    coords.append([center[0] + bounding_box[0][0], center[1] + bounding_box[1][1], center[2] + bounding_box[2][0]])
    coords.append([center[0] + bounding_box[0][0], center[1] + bounding_box[1][0], center[2] + bounding_box[2][1]])
    coords.append([center[0] + bounding_box[0][0], center[1] + bounding_box[1][0], center[2] + bounding_box[2][0]])
    return torch.tensor(coords)

def apply_transform_on_bb(bb_coords, q_mat):
    '''
    :param bb_coords: 8 x 3
    :param q_mat: 3 x 3
    :return: 8 x 3
    '''
    print(bb_coords)
    bb_coords = bb_coords.type(torch.float64)
    print(bb_coords)
    new_bb_coords = torch.matmul(bb_coords.type(torch.float64), torch.tensor(q_mat))
    return new_bb_coords

x_scale, y_scale, z_scale = 2, 1, 2
scale_mat = np.array([[x_scale, 0, 0],
                      [0, y_scale, 0],
                      [0, 0, z_scale]]).astype(np.float64)

# mesh = trimesh.load_mesh('chair10_m.obj')
mesh = IO.load_objs_as_meshes(['chair10_m.obj'])
# print(type(mesh))

# mesh_verts = mesh.verts_packed()
# print(mesh_verts.shape)
verts = mesh.verts_packed()
bb = mesh.get_bounding_boxes()
coords = get_bounding_box_coords(verts, bb[0])
coords = apply_transform_on_bb(coords, scale_mat)
print(coords)
# center = torch.mean(mesh_verts, dim=1)
# new_vert = torch.tensor([item.cpu().detach().numpy() for item in [i @ scale_mat for i in mesh.verts_padded()[0]]]).unsqueeze(0)
# new_vert = torch.tensor(new_vert)
# print(center)
# x = torch.tensor([1, 2, 3])
# # mat = matrix_to_quaternion(torch.tensor(scale_mat))
# # y = quaternion_apply(mat, x)
# y = x @ scale_mat
# print(y)

# print(center)
# print(mesh_verts.shape)
# for i in mesh_verts[0]:
#     print(i)
#     print(pytorch3d.transforms.quaternion_apply(torch.tensor(scale_mat), i))
# bb = mesh.get_bounding_boxes()
# print(bb)

# mesh = trimesh.Trimesh.apply_transform(mesh, scale_mat)
# bb = [[ 6.76065711e-03, -7.13646718e-03, -9.99951681e-01,
#          4.86275453e-04],
#        [ 9.99520434e-01,  3.02672878e-02,  6.54172950e-03,
#         -5.75688562e-03],
#        [ 3.02191405e-02, -9.99516364e-01,  7.33767152e-03,
#          9.71200397e-04],
#        [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,
#          1.00000000e+00]]
# bb_1 = [[ 1.62011877e-03,  9.99994615e-01, -2.85388357e-03,
#         -8.29902267e-04],
#        [-1.37179887e-03,  2.85610710e-03,  9.99994980e-01,
#         -1.95896088e-04],
#        [ 9.99997747e-01, -1.61619568e-03,  1.37641871e-03,
#          4.13472811e-04],
#        [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00,
#          1.00000000e+00]]

# bb_1, _ = trimesh.bounds.oriented_bounds(mesh)
# mesh = trimesh.Trimesh.apply_transform(mesh, bb)

# mesh.show()

# v = np.array(mesh.vertices)

# x_max = v[np.where(v[:, 0] == np.max(v[:, 0]))][0]
# x_min = v[np.where(v[:, 0] == np.min(v[:, 0]))][0]
# y_max = v[np.where(v[:, 1] == np.max(v[:, 1]))][0]
# y_min = v[np.where(v[:, 1] == np.min(v[:, 1]))][0]
# z_max = v[np.where(v[:, 2] == np.max(v[:, 2]))][0]
# z_min = v[np.where(v[:, 2] == np.min(v[:, 2]))][0]
# print(x_max)
# print(x_min)
# print(y_max)
# print(y_min)
# print(z_max)
# print(z_min)
# v_new = [x_max, x_min, y_max, y_min, z_max, z_min]
